# -*- coding: utf-8 -*-

# 去除单个单词构成的子图
import networkx as nx
import pickle
from bbdw.my_graph import Node, Edge, Graph


def remove():
    # 读取本地字典
    # 反序列化
    word_edge = []
    word_list = []

    with open(r'F:\mypython\final_subject\bbdw\exclude_doc_out\connected_graph.txt', 'rb') as f:
        graph = pickle.load(f)
        nodes_object = graph.nodes
        edges_object = graph.edges

        for edge in edges_object:
            word_edge.append([edge.node1.name, edge.node2.name])

        for node in nodes_object:
            word_list.append(node.name)

    point_list = word_list
    link_list = word_edge
    sub = []
    G = nx.Graph()
    for node1 in point_list:
        G.add_node(node1)
    for link in link_list:
        G.add_edge(link[0], link[1])

    num = 0
    for c in nx.connected_components(G):
        nodeSet = G.subgraph(c).nodes
        edgeSet = G.subgraph(c).edges
        if len(nodeSet) > 1:
            # print(nodeSet)
            # print(edgeSet)
            sub.append(nodeSet)
            num += 1
    print(num)

    left_word = []
    for node_view in sub:
        for word in node_view:
            left_word.append(word)

    new_nodes_obj = []
    new_edges_obj = []

    name_node_dict = {}
    for node in nodes_object:
        if node.name in left_word:
            new_node = Node(index=node.index, name=node.name, neighbors=node.neighbors, type=node.type)
            name_node_dict[node.name] = new_node
            new_nodes_obj.append(new_node)

    for node in new_nodes_obj:
        neighbors = node.neighbors
        new_neighbor = []
        for neighbor in neighbors:
            if neighbor in left_word:
                new_neighbor.append(neighbor)
        node.neighbors = new_neighbor

    for edge in edges_object:
        w1 = edge.node1.name
        w2 = edge.node2.name
        if w1 in left_word and w2 in left_word:
            node1 = name_node_dict[edge.node1.name]
            node2 = name_node_dict[edge.node2.name]
            weight = edge.weight
            new_edge = Edge(node1, node2, weight)
            new_edges_obj.append(new_edge)

    new_graph = Graph(new_nodes_obj, new_edges_obj)
    with open(r'F:\mypython\final_subject\bbdw\exclude_doc_single_word_out\connected_graph_exclude_doc.txt', 'wb') as f:
        pickle.dump(new_graph, f)


    # 加上文档节点
    with open(r'F:\mypython\final_subject\bbdw\exclude_doc_single_word_out\connected_graph_exclude_doc.txt', 'rb') as f:
        graph = pickle.load(f)

    with open(r'F:\mypython\final_subject\bbdw\out\doc_word_edge_list.txt', 'rb') as f:
        doc_word_edge_obj = pickle.load(f)

    name_node_dict = {}
    nodes_obj = graph.nodes  # 不包含文档节点
    word_word_edge = graph.edges   # 单词之间形成的边

    for node_obj in nodes_obj:
        name_node_dict[node_obj.name] = node_obj

    doc_nodes_obj = []
    new_doc_word_edge_obj = []

    name_doc_node_dict = {}

    for edge_obj in doc_word_edge_obj:
        doc_node = edge_obj.node1
        word_node = edge_obj.node2
        if word_node.name in left_word:
            if doc_node.name not in name_doc_node_dict:
                # neighbors = doc_node.neighbors
                # new_neighbor = []
                # for neighbor in neighbors:
                #     if neighbor in left_word:
                #         new_neighbor.append(neighbor)
                new_doc_node = Node(name=doc_node.name, index=doc_node.index,
                                    type=doc_node.type, neighbors=[])
                name_doc_node_dict[doc_node.name] = doc_node
                doc_nodes_obj.append(doc_node)

            doc_node = name_doc_node_dict[doc_node.name]

            # word_node = edge_obj.node2
            this_node = name_node_dict[word_node.name]
            this_node.add_neighbor(doc_node.name)
            doc_node.add_neighbor(this_node.name)

            new_doc_word_edge = Edge(doc_node, this_node, edge_obj.weight)
            new_doc_word_edge_obj.append(new_doc_word_edge)

    nodes_list = nodes_obj + doc_nodes_obj
    edge_list = word_word_edge + new_doc_word_edge_obj
    graph = Graph(nodes_list, edge_list)

    print(len(nodes_list))
    print(len(edge_list))

    for node in nodes_list:
        print(node.name)

    for edge in edge_list:
        print(edge.node1.name, edge.node2.name, edge.weight)

    with open(r'F:\mypython\final_subject\bbdw\include_doc_single_word_out\connected_graph.txt', 'wb') as f:
        pickle.dump(graph, f)


if __name__ == '__main__':
    remove()
